# 定义审批节点（判断是否通过）
from langgraph.graph import StateGraph, END
from langchain_core.runnables import Runnable
from graphviz import Source

from test.langgraph_base.ticket_node import intent_recognition_node, category_node, assign_node
from test.langgraph_base.ticket_state import TicketState


def approval_node(state: TicketState) -> TicketState:
    """审批工单，返回是否通过"""
    # 检查修改次数，如果已经修改过2次，无论如何都做最终处理
    if hasattr(state, 'modification_count') and state.modification_count >= 2:
        is_approved = False  # 最终拒绝
        is_final_rejection = True
        state.history.append({"node": "approval", "result": "最终拒绝（达到最大修改次数）"})
    else:
        is_approved = state.category != "其他"  # 示例：非"其他"分类自动通过
        is_final_rejection = False
        state.history.append({"node": "approval", "result": "通过" if is_approved else "驳回"})
    
    return state.model_copy(update={
        "is_approved": is_approved,
        "is_final_rejection": is_final_rejection
    })

# 定义修改节点（审批驳回时执行）
def modify_node(state: TicketState) -> TicketState:
    """修改工单内容（示例：补充说明）"""
    modified_query = state.user_query + " [补充：请明确产品型号]"
    state.history.append({"node": "modify", "result": modified_query})
    # 增加修改次数计数
    new_count = getattr(state, 'modification_count', 0) + 1
    return state.model_copy(update={
        "user_query": modified_query,
        "modification_count": new_count
    })

# 构建带条件分支的图
workflow = StateGraph(TicketState)

# 1. 添加所有节点
workflow.add_node("intent_recognition", intent_recognition_node)
workflow.add_node("category", category_node)
workflow.add_node("approval", approval_node)
workflow.add_node("modify", modify_node)
workflow.add_node("assign", assign_node)

# 2. 定义条件分支路由函数
def approval_router(state: TicketState) -> str:
    """路由函数：返回下一步节点名"""
    if state.is_approved:
        return "assign"  # 通过→分配
    elif hasattr(state, 'is_final_rejection') and state.is_final_rejection:
        return END  # 最终拒绝→结束
    else:
        return "modify"  # 驳回→修改

# 3. 设置工作流流程
workflow.set_entry_point("intent_recognition")  # 设置入口点
workflow.add_edge("intent_recognition", "category")  # 意图识别→分类
workflow.add_edge("category", "approval")  # 分类→审批
workflow.add_conditional_edges("approval", approval_router)  # 审批→条件分支
workflow.add_edge("modify", "intent_recognition")  # 修改后→意图识别（循环）
workflow.add_edge("assign", END)  # 分配→结束

app = workflow.compile()

graph = app.get_graph()

# 构建 DOT 语言字符串
dot_str = "digraph G {\n"
# 添加节点
for node in graph.nodes:
    dot_str += f'  "{node}" [label="{node}"];\n'
# 添加边（包括普通边和条件边）
for edge in graph.edges:
    start = edge.source
    end = edge.target
    # 条件边可添加标签（如路由函数名）
    label = edge.condition.__name__ if hasattr(edge, 'condition') and edge.condition else ""
    dot_str += f'  "{start}" -> "{end}" [label="{label}"];\n'
dot_str += "}"



# 渲染并保存
Source(dot_str).render("workflow", format="png", cleanup=True)
# 获取图表数据
# mermaid_output = graph.draw_mermaid_png()
#
# # 处理返回的数据 - draw_mermaid_png()返回PNG二进制数据
# if isinstance(mermaid_output, bytes):
#     # 直接以二进制模式保存PNG图片
#     with open("workflow.png", "wb") as f:
#         f.write(mermaid_output)
#     print("\n工作流图表已保存为PNG图片: workflow.png")
# else:
#     # 如果是文本（不常见情况），则保存为文本
#     try:
#         with open("workflow_mermaid.txt", "w", encoding="utf-8") as f:
#             f.write(str(mermaid_output))
#         print("\nMermaid图表代码已保存到 workflow_mermaid.txt")
#     except Exception as e:
#         print(f"\n保存图表失败: {e}")
#
# print("\n提示: 已完成工作流可视化处理")

if __name__ == '__main__':
    print("===== 工单审批流程测试 =====")
    
    # 测试场景1：审批通过的工单（非"其他"分类）
    print("\n[测试场景1] 审批通过的工单")
    initial_state1 = TicketState(ticket_id="APPROVAL-001", user_query="我想退款，产品质量有问题")
    print(f"初始状态: {initial_state1.model_dump()}")
    
    final_state1 = app.invoke(initial_state1)
    
    print(f"\n处理完成!")
    print(f"工单ID: {final_state1.get('ticket_id', 'N/A') if isinstance(final_state1, dict) else getattr(final_state1, 'ticket_id', 'N/A')}")
    print(f"用户请求: {final_state1.get('user_query', 'N/A') if isinstance(final_state1, dict) else getattr(final_state1, 'user_query', 'N/A')}")
    print(f"识别意图: {final_state1.get('intent', 'N/A') if isinstance(final_state1, dict) else getattr(final_state1, 'intent', 'N/A')}")
    print(f"工单分类: {final_state1.get('category', 'N/A') if isinstance(final_state1, dict) else getattr(final_state1, 'category', 'N/A')}")
    print(f"审批状态: {final_state1.get('is_approved', 'N/A') if isinstance(final_state1, dict) else getattr(final_state1, 'is_approved', 'N/A')}")
    print(f"分配人员: {final_state1.get('assigned_to', 'N/A') if isinstance(final_state1, dict) else getattr(final_state1, 'assigned_to', 'N/A')}")
    
    print(f"\n执行历史:")
    history1 = final_state1.get('history', []) if isinstance(final_state1, dict) else getattr(final_state1, 'history', [])
    for step in history1:
        print(f"  - {step.get('node', 'N/A')}: {step.get('result', 'N/A')}")
    
    # 测试场景2：审批驳回并修改的工单（"其他"分类）
    print("\n\n[测试场景2] 审批驳回并修改的工单")
    initial_state2 = TicketState(ticket_id="APPROVAL-002", user_query="我有一个建议")
    print(f"初始状态: {initial_state2.model_dump()}")
    
    # 执行第二个测试场景
    final_state2 = app.invoke(initial_state2)
    
    print(f"\n处理完成!")
    print(f"工单ID: {final_state2.get('ticket_id', 'N/A') if isinstance(final_state2, dict) else getattr(final_state2, 'ticket_id', 'N/A')}")
    print(f"用户请求: {final_state2.get('user_query', 'N/A') if isinstance(final_state2, dict) else getattr(final_state2, 'user_query', 'N/A')}")
    print(f"识别意图: {final_state2.get('intent', 'N/A') if isinstance(final_state2, dict) else getattr(final_state2, 'intent', 'N/A')}")
    print(f"工单分类: {final_state2.get('category', 'N/A') if isinstance(final_state2, dict) else getattr(final_state2, 'category', 'N/A')}")
    print(f"审批状态: {final_state2.get('is_approved', 'N/A') if isinstance(final_state2, dict) else getattr(final_state2, 'is_approved', 'N/A')}")
    print(f"分配人员: {final_state2.get('assigned_to', 'N/A') if isinstance(final_state2, dict) else getattr(final_state2, 'assigned_to', 'N/A')}")
    
    print(f"\n执行历史:")
    history2 = final_state2.get('history', []) if isinstance(final_state2, dict) else getattr(final_state2, 'history', [])
    for step in history2:
        print(f"  - {step.get('node', 'N/A')}: {step.get('result', 'N/A')}")
    
    print("\n===== 测试完成 =====")
    print("工单审批流程已成功执行。")